import os
import uuid
import pickle
import cv2
import numpy as np
from flask import Flask, render_template, request, jsonify, send_from_directory
import face_recognition

app = Flask(__name__)

# 文件上传的路径
UPLOAD_FOLDER = 'uploads'
RESULT_FOLDER = 'results'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['RESULT_FOLDER'] = RESULT_FOLDER
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)
if not os.path.exists(RESULT_FOLDER):
    os.makedirs(RESULT_FOLDER)

# 初始化人脸库（字典）
face_db = {}

# 加载已保存的face_db（如果有）
def load_face_db(filepath='face_database.pkl'):
    global face_db
    if os.path.exists(filepath):
        with open(filepath, 'rb') as f:
            face_db = pickle.load(f)

# 保存face_db
def save_face_db(filepath='face_database.pkl'):
    with open(filepath, 'wb') as f:
        pickle.dump(face_db, f)

# 添加人脸到人脸库
def add_face_to_library(image_path, name):
    image = face_recognition.load_image_file(image_path)
    face_encodings = face_recognition.face_encodings(image)

    if len(face_encodings) > 0:
        face_encoding = face_encodings[0]
        face_id = str(uuid.uuid4())
        face_db[face_id] = {'name': name, 'encoding': face_encoding}
        save_face_db()
        return face_id, name
    return None, None

# 编辑人脸信息
def edit_face_info(face_id, new_name):
    if face_id in face_db:
        face_db[face_id]['name'] = new_name
        save_face_db()
        return True
    return False

# 删除人脸记录
def delete_face(face_id):
    if face_id in face_db:
        del face_db[face_id]
        save_face_db()
        return True
    return False

# 比较上传的图片中的人脸
def compare_faces(image_path):
    image = face_recognition.load_image_file(image_path)
    face_locations = face_recognition.face_locations(image)
    face_encodings = face_recognition.face_encodings(image)

    img = cv2.imread(image_path)
    if img is None:
        print(f"Error: Unable to load image from path: {image_path}")
        return None, []

    result_image_path = os.path.join(app.config['RESULT_FOLDER'], os.path.basename(image_path))
    results = []

    if len(face_encodings) == 0:
        return None, [{"error": "No faces detected in the uploaded image."}]

    for (top, right, bottom, left), face_encoding in zip(face_locations, face_encodings):
        matches = face_recognition.compare_faces([face['encoding'] for face in face_db.values()], face_encoding)
        name = "Unknown"
        face_id = "Unknown"

        if True in matches:
            match_index = matches.index(True)
            matched_face = list(face_db.items())[match_index]
            face_id, info = matched_face
            name = info['name']

        cv2.rectangle(img, (left, top), (right, bottom), (0, 255, 0), 2)
        cv2.putText(img, f"{name} (ID: {face_id})", (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2)

        results.append({"name": name, "id": face_id, "location": (top, right, bottom, left)})

    cv2.imwrite(result_image_path, img)
    return result_image_path, results

@app.route('/')
def index():
    return render_template('index.html', faces=face_db)

@app.route('/upload_face', methods=['POST'])
def upload_face():
    if 'face_image' not in request.files or 'name' not in request.form:
        return jsonify({"success": False, "message": "Invalid input"}), 400

    file = request.files['face_image']
    name = request.form['name']

    if file.filename == '':
        return jsonify({"success": False, "message": "No file uploaded"}), 400

    face_path = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
    file.save(face_path)

    face_id, added_name = add_face_to_library(face_path, name)

    if face_id:
        return jsonify({"success": True, "id": face_id, "name": added_name})
    else:
        return jsonify({"success": False, "message": "No face found, please try again."}), 400

@app.route('/edit_face', methods=['POST'])
def edit_face():
    face_id = request.form.get('face_id')
    new_name = request.form.get('new_name')
    if edit_face_info(face_id, new_name):
        return jsonify({"success": True, "message": "Face info updated successfully."})
    else:
        return jsonify({"success": False, "message": "Face ID not found."}), 400

@app.route('/delete_face', methods=['POST'])
def delete_face_route():
    face_id = request.form.get('face_id')
    if delete_face(face_id):
        return jsonify({"success": True, "message": "Face deleted successfully."})
    else:
        return jsonify({"success": False, "message": "Face ID not found."}), 400

@app.route('/compare', methods=['POST'])
def compare():
    if 'compare_image' not in request.files:
        return jsonify({"success": False, "message": "No file uploaded"}), 400

    file = request.files['compare_image']

    if file.filename == '':
        return jsonify({"success": False, "message": "No file uploaded"}), 400

    compare_path = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
    file.save(compare_path)

    result_image_path, results = compare_faces(compare_path)

    if results and "error" in results[0]:
        return jsonify({"success": False, "message": results[0]["error"]})

    return jsonify({"success": True, "results": results, "result_image": os.path.basename(result_image_path)})

@app.route('/results/<filename>')
def show_result(filename):
    return send_from_directory(app.config['RESULT_FOLDER'], filename)

if __name__ == '__main__':
    load_face_db()
    app.run(debug=True)
